package cn.ymatrix.mxgateKafkaSDKTOYMatrix;

import cn.ymatrix.apiclient.DataPostListener;
import cn.ymatrix.apiclient.MxClient;
import cn.ymatrix.apiclient.Result;
import cn.ymatrix.builder.MxBuilder;
import cn.ymatrix.builder.RequestType;
import cn.ymatrix.concurrencycontrol.WorkerPool;
import cn.ymatrix.concurrencycontrol.WorkerPoolFactory;
import cn.ymatrix.data.Tuple;
import cn.ymatrix.data.Tuples;
import cn.ymatrix.exception.AllTuplesFailException;
import cn.ymatrix.exception.PartiallyTuplesFailException;
import cn.ymatrix.logger.LoggerLevel;
import cn.ymatrix.logger.MxLogger;
import cn.ymatrix.utils.Config;
import com.alibaba.fastjson.JSON;
import org.slf4j.Logger;

import java.util.List;
import java.util.Map;

public class MxgateConsumerKafkaSDKTOYMatrix {

  private static final String CUSTOMER_LOG_TAG = "[>>>CUSTOMER<<<] ";
  private static MxBuilder mxBuilder;

  public static void main(String[] args) {
    MxLogger.loggerLevel(LoggerLevel.INFO);
    Logger LOGGER = MxLogger.init(MxgateConsumerKafkaSDKTOYMatrix.class);

    initMxgateSDKBuilder(LOGGER);

    mxgateMxClientBuilder(LOGGER);
  }

  private static void initMxgateSDKBuilder(Logger LOGGER) {
    String requestType = Config.getRequestType();
    String dropAll = Config.getDropAll();
    String asyncMode = Config.getAsyncMode();
    MxLogger.loggerLevel(LoggerLevel.INFO);

    MxBuilder.Builder builder =
        MxBuilder.newBuilder()
            .withDropAll(false)
            .withCacheCapacity(Config.getCacheCapacity())
            .withCacheEnqueueTimeout(Config.getCacheEnqueueTimeout())
            .withConcurrency(Config.getSdkConcurrency())
            .withRequestTimeoutMillis(Config.getRequestTimeoutMillis())
            .withMaxRequestQueued(Config.getMaxRequestQueued())
            .withMaxRetryAttempts(Config.getMaxRetryAttempts())
            .withRetryWaitDurationMillis(Config.getRetryWaitDurationMillis());

    if (requestType.equals("http")) {
      builder.withRequestType(RequestType.WithHTTP);
    } else if (requestType.equals("grpc")) {
      builder.withRequestType(RequestType.WithGRPC);
    } else {
      LOGGER.error("Invalid request type, should be http or grpc");
      return;
    }
    if (dropAll.equals("dropAll")) {
      builder.withDropAll(true);
    }
    if (asyncMode.equals("async")) {
      builder.withRequestAsync(true);
    }
    mxBuilder = builder.build();
  }

  public static void mxgateMxClientBuilder(Logger LOGGER) {

    String linesOrBytesToFlush = Config.getLinesOrBytesToFlush();
    int enoughLinesOrBytes = Config.getEnoughLinesOrBytes();
    int csvConstructionBatchSize = Config.getCsvConstructionBatchSize();
    String compressWithZSTD = Config.getCompressWithZstd();
    String requestType = Config.getRequestType();

    String httpHost = Config.getHttpUrl();
    String gRPCHost = Config.getgRPCHost();
    String schema = Config.getSchemaName();
    String table = Config.getTableName();

    int consumersSize = Config.getConsumersSize();
    MxClient client = null;
    try {
      if (requestType.equals("http")) {
        client = mxBuilder.connect(httpHost, gRPCHost, schema, table);
      } else if (requestType.equals("grpc")) {
        client = mxBuilder.connect(gRPCHost, gRPCHost, schema, table);
      }
      if (client != null) {
        // client.withoutCompress();
        if (csvConstructionBatchSize > 0) {
          client.withCSVConstructionBatchSize(csvConstructionBatchSize);
        }
        if (linesOrBytesToFlush.equals("lines")) {
          // NOTE: 使用按行 flush 的速度要比按 bytes 快。
          client.withEnoughLinesToFlush(enoughLinesOrBytes);
        } else if (linesOrBytesToFlush.equals("bytes")) {
          client.withEnoughBytesToFlush(enoughLinesOrBytes);
        } else {
          LOGGER.error("Invalid flush mode :" + linesOrBytesToFlush);
          return;
        }
        if (compressWithZSTD.equals("zstd")) {
          LOGGER.info("with compress");
          client.withCompress();
          if (requestType.equals("grpc")) {
            LOGGER.info("with base64");
            client.withBase64Encode4Compress();
          }
        }
        client.registerDataPostListener(
            new DataPostListener() {
              public void onSuccess(Result result) {
                // 获取offset值
                Tuples tuples = result.getRawTuples();
                List<Tuple> tupleList = tuples.getTuplesList();
                for (Tuple tuple : tupleList) {
                  tuple.getSerialNum();
                }
              }

              public void onFailure(Result result) {
                LOGGER.error(
                    CUSTOMER_LOG_TAG
                        + "Send tuples fail msg: {} lines: {}"
                        + result.getMsg()
                        + result.getSucceedLines());
                LOGGER.error(
                    CUSTOMER_LOG_TAG
                        + "Send tuples fail error tuples: {}"
                        + result.getErrorTuplesMap());
                for (Map.Entry<Tuple, String> entry : result.getErrorTuplesMap().entrySet()) {
                  LOGGER.error(
                      CUSTOMER_LOG_TAG
                          + "error tuple of table="
                          + entry.getKey().getTableName()
                          + "tuple="
                          + entry.getKey()
                          + "reason="
                          + entry.getValue());
                }
              }
            });
      }

      // 开启程序开始消费
      consumeKafkaTOYMatrix(consumersSize, LOGGER, client);
      Runtime.getRuntime()
          .addShutdownHook(
              new Thread(
                  () -> {
                    LOGGER.error("Thread down ..." + mxBuilder.getTupleCacheSize());
                  }));
    } catch (Exception e) {
      LOGGER.error("MxClient init error: {}" + e.getMessage());
      e.printStackTrace();
    }
  }

  public static void consumeKafkaTOYMatrix(int consumersSize, Logger LOGGER, MxClient client) {

    String kafkaBootStrap = Config.getKafkaBootStrap();
    String kafkaTopic = Config.getKafkaTopic();
    String kafkaConsumerGroup = Config.getKafkaConsumerGroup();
    int timeoutMillis = Config.getTimeoutMillis();

    WorkerPool consumerPool = WorkerPoolFactory.initFixedSizeWorkerPool(consumersSize);
    for (int i = 0; i < consumersSize; i++) {
      int finalI = i;
      Runnable runnable =
          new Runnable() {
            public void run() {
              KafkaConsumerWrapper consumer =
                  new KafkaConsumerWrapper(kafkaBootStrap, kafkaTopic, finalI, kafkaConsumerGroup);
              consumer.consumeLoop(
                  timeoutMillis,
                  new KafkaMsgConsumeCallback() {
                    public void convertMsgValueToObj(String value, int num) {
                      try {

                        UserTable tt = JSON.parseObject(value, UserTable.class);
                        Tuple tuple = client.generateEmptyTuple();
                        tuple.addColumn("ts", tt.getTs());
                        tuple.addColumn("tag", tt.getTag());
                        tuple.addColumn("c1", tt.getC1());
                        tuple.addColumn("c2", tt.getC2());
                        tuple.addColumn("c3", tt.getC3());
                        tuple.addColumn("c4", tt.getC4());
                        tuple.addColumn("c5", tt.getC5());
                        tuple.addColumn("c6", tt.getC6());
                        tuple.addColumn("c7", tt.getC7());
                        tuple.addColumn("c8", tt.getC8());

                        if (client.appendTupleBlocking(tuple)) {
                          long flushStartTime = System.currentTimeMillis();
                          client.flushBlocking();
                          LOGGER.info(
                              "--> flush data to YMatrix COST: "
                                  + (System.currentTimeMillis() - flushStartTime)
                                  + " ms");
                        }

                      } catch (AllTuplesFailException e) {
                        Result result = e.getResult();
                        for (Map.Entry<Tuple, String> entry :
                            result.getErrorTuplesMap().entrySet()) {
                          LOGGER.error(
                              "error tuple of table={}, tuple={}, reason={}",
                              entry.getKey().getTableName(),
                              entry.getKey(),
                              entry.getValue());
                        }
                      } catch (PartiallyTuplesFailException e) {
                        Result result = e.getResult();
                        for (Map.Entry<Tuple, String> entry :
                            result.getErrorTuplesMap().entrySet()) {
                          LOGGER.error(
                              "error tuple of table={}, tuple={}, reason={}",
                              entry.getKey().getTableName(),
                              entry.getKey(),
                              entry.getValue());
                        }
                      } catch (Exception e) {
                        LOGGER.error("Exception error, info={}", e);
                      }
                    }
                  });
            }
          };
      consumerPool.join(runnable);
    }
  }
}
